package main

import (
	"fmt"
	"log"
)

// Hub 中心枢纽
type Hub struct {
	clients    map[*Client]bool            // 所有在线客户端
	userMap    map[string]map[*Client]bool // 用户ID，该用户所有连接的客户端，因为是多端登录，所以需要用map存储
	rooms      map[string]map[*Client]bool // 房间ID，该房间所有连接的客户端
	register   chan *Client                // 注册新客户端的通道
	unregister chan *Client                // 注销客户端的通道
	broadcast  chan Broadcast              // 广播消息的通道
}

func NewHub() *Hub {
	return &Hub{
		clients:    make(map[*Client]bool),
		userMap:    make(map[string]map[*Client]bool),
		rooms:      make(map[string]map[*Client]bool),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		broadcast:  make(chan Broadcast),
	}
}

func (h *Hub) run() {
	for {
		select {
		case c := <-h.register:
			h.clients[c] = true // 加入客户端列表
			if _, ok := h.userMap[c.user]; !ok {
				h.userMap[c.user] = make(map[*Client]bool) // 如果用户ID不存在，则创建
			}
			h.userMap[c.user][c] = true // 加入用户列表
			log.Println("New client connected for user:", c.user)
		case c := <-h.unregister:
			if _, ok := h.clients[c]; ok {
				delete(h.clients, c) // 从客户端列表中移除
				if conn, ok := h.userMap[c.user]; ok {
					delete(conn, c) // 从用户列表中移除
					if len(conn) == 0 {
						delete(h.userMap, c.user) // 如果该用户没有其他连接，则从用户列表中移除
					}
				}
			}
			// 删除房间
			for room := range h.rooms { // 遍历客户端加入的所有房间
				if clients, ok := h.rooms[room]; ok { // 如果房间存在
					delete(clients, c) // 从房间中移除客户端
					if len(clients) == 0 {
						delete(h.rooms, room) // 如果房间中没有其他客户端，则从房间列表中移除
					}
				}
			}
			close(c.send)
			fmt.Printf("Client disconnected: %s\n", c.user)
		case b := <-h.broadcast:
			// 判断b.room 群聊
			if b.Room != "" {
				if clients, ok := h.rooms[b.Room]; ok { // 如果房间存在
					for client := range clients {
						select {
						case client.send <- b.Data: // 向room的所有客户端发送数据
						default: // 当有慢连接或者假连接，它将被关闭
							go func(c *Client) {
								h.unregister <- c
								c.conn.Close()
							}(client)
						}
					}
				}
			} else if b.To != "" { // 判断b.To 私聊
				if conns, ok := h.userMap[b.To]; ok { // 如果用户存在
					for client := range conns {
						select {
						case client.send <- b.Data: // 向该用户的所有客户端(多端)发送数据
						default: // 当有慢连接或者假连接，它将被关闭
							go func(c *Client) {
								h.unregister <- c
								c.conn.Close()
							}(client)
						}
					}
				}
			} else { // 默认广播
				for client := range h.clients {
					select {
					case client.send <- b.Data: // 向所有客户端发送数据
					default: // 当有慢连接或者假连接，它将被关闭
						go func(c *Client) {
							h.unregister <- c
							c.conn.Close()
						}(client)
					}
				}
			}
		}
	}
}
